import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
import logging

logger = logging.getLogger(__name__)

def ensure_datetime_column(df, column='timestamp', timezone='Asia/Shanghai'):
    """
    确保DataFrame中的日期时间列是datetime类型并带有指定时区。
    
    参数:
        df: pandas DataFrame
        column: 日期时间列名
        timezone: 目标时区
        
    返回:
        处理后的DataFrame副本
    """
    from source.data_processor.utils.time_utils import standardize_to_shanghai
    
    if df.empty or column not in df.columns:
        return df.copy()
    
    result_df = df.copy()
    
    # 转换为datetime类型
    if not pd.api.types.is_datetime64_any_dtype(result_df[column]):
        result_df[column] = pd.to_datetime(result_df[column])
    
    # 处理时区
    if pd.api.types.is_datetime64tz_dtype(result_df[column]):
        # 已有时区信息，转换到目标时区
        result_df[column] = result_df[column].dt.tz_convert(timezone)
    else:
        # 无时区信息，假设为UTC并转换
        result_df[column] = result_df[column].dt.tz_localize('UTC').dt.tz_convert(timezone)
    
    return result_df

def filter_dataframe_by_time(df, time_point, column='timestamp', before=True):
    """
    根据时间点筛选DataFrame。
    
    参数:
        df: pandas DataFrame
        time_point: 时间点（datetime对象或能被pd.to_datetime解析的字符串）
        column: 日期时间列名
        before: 如果为True，保留小于时间点的数据；如果为False，保留大于等于时间点的数据
        
    返回:
        筛选后的DataFrame
    """
    from data_processor.utils.time_utils import standardize_to_shanghai
    
    if df.empty or column not in df.columns:
        return df.copy()
    
    # 确保df[column]是datetime类型
    df = ensure_datetime_column(df, column)
    
    # 确保time_point是datetime类型并且具有时区信息
    if not isinstance(time_point, pd.Timestamp):
        time_point = pd.to_datetime(time_point)
    
    time_point = standardize_to_shanghai(time_point)
    
    # 筛选数据
    if before:
        return df[df[column] < time_point].copy()
    else:
        return df[df[column] >= time_point].copy()

def smooth_timeseries(data_series, method='linear', limit=5):
    """
    对时间序列数据进行平滑处理，填充缺失值。
    
    参数:
        data_series: 需要平滑的数据Series
        method: 平滑方法，可选值: 'linear', 'ffill', 'bfill', 'nearest'等，详见pandas.interpolate
        limit: 连续填充的最大缺失值数量
        
    返回:
        平滑处理后的数据Series
    """
    if data_series.empty:
        return data_series
    
    # 复制原始数据，避免修改原始对象
    smoothed_series = data_series.copy()
    
    # 检查NaN值数量
    nan_count = smoothed_series.isna().sum()
    if nan_count > 0:
        logger.debug(f"检测到 {nan_count} 个缺失值，进行数据平滑")
        
        try:
            # 对于大型缺口，先尝试使用前向填充处理
            if nan_count > limit:
                # 先使用前向填充处理大型缺口的开始部分
                smoothed_series.ffill(inplace=True, limit=1)
                # 再使用后向填充处理大型缺口的结束部分
                smoothed_series.bfill(inplace=True, limit=1)
            
            # 然后使用插值处理剩余缺失值
            smoothed_series = smoothed_series.interpolate(
                method=method,
                limit=limit,
                limit_direction='both'
            )
            
            # 如果仍有缺失值（可能在序列开头或结尾），尝试前向和后向填充
            if smoothed_series.isna().sum() > 0:
                smoothed_series.ffill(inplace=True)
                smoothed_series.bfill(inplace=True)
                
            remaining_nan = smoothed_series.isna().sum()
            if remaining_nan > 0:
                logger.warning(f"数据平滑后仍有 {remaining_nan} 个缺失值")
            else:
                logger.debug("数据平滑成功完成")
                
        except Exception as e:
            logger.error(f"数据平滑处理时发生错误: {e}")
    
    return smoothed_series

def standardize_data(data_series):
    """
    对数据进行标准化处理，同时保留原始值。
    
    参数:
        data_series: 需要标准化的数据Series
        
    返回:
        包含标准化值和原始值的元组 (standardized_values, original_values)
    """
    if data_series.empty:
        logger.warning("输入数据为空，无法进行标准化")
        return pd.Series(dtype=float), data_series
    
    # 保存原始值
    original_values = data_series.copy()
    
    # 先检查是否全为相同值
    if data_series.nunique() == 1:
        logger.debug(f"数据全为同一值 ({data_series.iloc[0]})，标准化结果将全为0")
        # 直接创建全0序列
        standardized_series = pd.Series(0.0, index=data_series.index)
        return standardized_series, original_values
    
    # 进行数据平滑处理，填充缺失值
    smoothed_series = smooth_timeseries(data_series)
    
    # 创建标准化后的Series，初始全为NaN
    standardized_series = pd.Series(index=data_series.index, dtype=float)
    
    try:
        # 获取非NaN数据和索引
        non_na_mask = ~smoothed_series.isna()
        valid_indices = smoothed_series.index[non_na_mask]
        valid_values = smoothed_series.values[non_na_mask]
        
        if len(valid_values) == 0:
            logger.warning("平滑处理后无有效数据，无法标准化")
            return pd.Series(0.0, index=data_series.index), original_values
        
        # 处理无方差或只有一个值的情况
        if len(valid_values) == 1 or np.std(valid_values) < 1e-8:
            logger.debug("数据方差接近0，标准化结果将全为0")
            standardized_series[valid_indices] = 0.0
        else:
            # 标准化处理
            scaler = StandardScaler()
            standardized_values = scaler.fit_transform(valid_values.reshape(-1, 1)).flatten()
            
            # 确保长度一致，这里使用索引布尔掩码直接赋值，避免索引不匹配问题
            standardized_series[valid_indices] = standardized_values
        
        # 如果还有NaN值，对标准化后的序列再次进行平滑处理
        if standardized_series.isna().sum() > 0:
            standardized_series = smooth_timeseries(standardized_series)
        
        return standardized_series, original_values
    
    except Exception as e:
        logger.error(f"标准化数据时发生错误: {e}")
        # 出错时返回零值序列
        return pd.Series(0.0, index=data_series.index), original_values 